package com.dliyun.fort.gateway.ssh;

/**
 * @author jtoms.shen
 * @version 1.0
 * @date 2019/3/21 10:39
 */

import com.jcraft.jsch.ChannelShell;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.Session;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.io.OutputStream;

@Slf4j
public class SSHClient {

    private ChannelShell channel;

    private Session session;

    public SSHClient() {
    }

    public static SSHClient connect(String host, int port, String userName, String password, int width, int height) throws Exception {
        JSch jsch = new JSch();
        Session session = jsch.getSession(userName, host, port);

        session.setConfig("StrictHostKeyChecking", "no");
        session.setPassword(password);
        session.connect();
        ChannelShell channel = (ChannelShell) session.openChannel("shell");
        channel.connect(3000);
        channel.setPty(true);
        channel.setPtySize(width, height, 0, 0);


        SSHClient client = new SSHClient();
        client.session = session;
        client.channel = channel;
        return client;
    }

    public static SSHClient connectWithSSHKey(String host, int port, String user, String privateKey, String password, int width, int height) throws Exception {

        JSch jsch = new JSch();
        jsch.addIdentity("", privateKey.getBytes(), null, password.getBytes());

        Session session = jsch.getSession(user, host, port);
        session.setConfig("StrictHostKeyChecking", "no");
        session.connect();

        ChannelShell channel = (ChannelShell) session.openChannel("shell");
        channel.connect(3000);
        channel.setPty(true);
        channel.setPtySize(width, height, 0, 0);

        SSHClient client = new SSHClient();
        client.session = session;
        client.channel = channel;
        return client;
    }

    public void setWindow(int width, int height) throws IOException {
        this.channel.setPtySize(width, height, 0, 0);
    }

    public void write(String cmd) throws IOException {
        try {
            if (channel != null) {
                OutputStream outputStream = channel.getOutputStream();
                outputStream.write(cmd.getBytes());
                outputStream.flush();
            }
        } catch (Exception e) {
            e.printStackTrace();
            throw e;
        }
    }

    public ChannelShell getChannel() {
        return this.channel;
    }

    public Session getSession() {
        return this.session;
    }

    public void close() {
        try {
            this.channel.disconnect();
            this.session.disconnect();
        } catch (Exception e) {
            log.error("", e);
        }
    }
}

